import numpy as np
import torch
import torch.nn.functional as F

from attack.base_attack import BaseAttack

class Attack(BaseAttack):
    def __init__(self, model, config):
        super().__init__('TMI+FGSM', model, config, batch_size = 100, targeted = True, llc = False)

    def attack_batch(self, images, targets):
        
        adv_images = np.copy(images)
        momentum = 0

        for num_steps in range(self.config['num_steps']):
            var_images = torch.from_numpy(adv_images).to(self.model.device)
            var_images.requires_grad = True
            var_targets = torch.LongTensor(targets).to(self.model.device)

            self.model.eval()

            output = self.model(var_images)
            loss = F.cross_entropy(output, var_targets)
            loss.backward()
            grad = var_images.grad.cpu().numpy()
            momentum = self.config['decay_factor'] * momentum + grad

            adv_images = adv_images - self.config['eps_iter'] * np.sign(momentum)
            adv_images = np.clip(adv_images, images - self.config['epsilon'], images + self.config['epsilon'])
            adv_images = np.clip(adv_images, 0.0, 1.0)

            var_images.grad.data.zero_()

        return adv_images
